#include "yaml-cpp/yaml.h"
#include <iostream>
#include <iostream>
#include "opengauss/libpq-fe.h"
#include "file_utils.h"
#include <fstream>
#include <cstring> // 用于strcpy

extern void check_disk_space(const char *path, double *use_percent);

std::string getsource(const char *location)
{
    YAML::Node config = YAML::LoadFile(location);
    std::string datasource;
    datasource.append("dbname=");
    datasource.append(config["gaussdb"]["source"]["database"].as<std::string>());
    datasource.append(" port=");
    datasource.append(config["gaussdb"]["source"]["port"].as<std::string>());
    datasource.append(" application_name=libpq connect_timeout=5 sslmode=allow");
    datasource.append(" user=");
    datasource.append(config["gaussdb"]["source"]["username"].as<std::string>());
    return datasource;
}

// 单独的创建PGconn连接的函数
PGconn *createPGConnection(const std::string &connectionInfo)
{
    PGconn *conn = PQconnectdb(connectionInfo.c_str());
    if (PQstatus(conn) != CONNECTION_OK)
    {
        std::cerr << "连接数据库出错: " << PQerrorMessage(conn) << std::endl;
        PQfinish(conn);
        return nullptr;
    }
    return conn;
}

int GetActivityPids(PGconn *conn, const char *sql, std::vector<std::string> &pid)
{
    int nFields;
    int i, j;
    // 执行查询语句
    PGresult *res = PQexec(conn, sql);
    if (PQresultStatus(res) != PGRES_TUPLES_OK)
    {
        fprintf(stderr, "查询失败: %s", PQerrorMessage(conn));
        PQclear(res);
        PQfinish(conn);
        return 1;
    }
    /* 打印属性名称 */
    nFields = PQnfields(res);
    /* 打印行 */
    for (i = 0; i < PQntuples(res); i++)
    {
        for (j = 0; j < nFields; j++)
        {
            pid.push_back(PQgetvalue(res, i, j));
        }
    }
    // 释放结果集和关闭数据库连接
    PQclear(res);
    return 0;
}

// 终止指定PID的后端进程
bool PGTerminateBackend(PGconn *conn, std::string pid)
{
    if (!conn || PQstatus(conn) != CONNECTION_OK)
    {
        std::cerr << "错误：无效的数据库连接" << std::endl;
        return false;
    }

    // 构建终止进程的SQL命令
    std::string sql = "SELECT pg_terminate_backend(" + pid + ");";

    // 执行SQL
    PGresult *res = PQexec(conn, sql.c_str());
    if (!res)
    {
        std::cerr << "执行失败：内存分配错误" << std::endl;
        return false;
    }

    // 检查执行结果
    bool success = false;
    ExecStatusType status = PQresultStatus(res);

    if (status == PGRES_TUPLES_OK)
    {
        // pg_terminate_backend返回布尔值(t/f)，检查第一行第一列的结果
        if (PQntuples(res) > 0 && PQnfields(res) > 0)
        {
            const char *result = PQgetvalue(res, 0, 0);
            success = (result && std::string(result) == "t");
        }
    }
    else
    {
        std::cerr << "终止进程失败：" << PQerrorMessage(conn) << std::endl;
    }

    // 释放资源
    PQclear(res);
    return success;
}
// 执行查询并通过指针参数返回data_directory路径
// 返回值：0表示成功，非0表示失败
int GetPGDataDirectory(PGconn *conn, char *DataDirectory, size_t bufferSize)
{
    // 检查参数有效性
    if (!conn || PQstatus(conn) != CONNECTION_OK)
    {
        std::cerr << "错误：数据库连接无效或未建立" << std::endl;
        return 1;
    }
    if (!DataDirectory || bufferSize == 0)
    {
        std::cerr << "错误：输出缓冲区无效" << std::endl;
        return 2;
    }

    // 执行查询语句
    const char *sql = "SELECT setting FROM pg_settings WHERE name = 'data_directory';";
    PGresult *res = PQexec(conn, sql);

    if (!res)
    {
        std::cerr << "查询失败：内存分配错误" << std::endl;
        return 3;
    }

    // 检查查询结果状态
    int result = 4; // 默认失败
    ExecStatusType status = PQresultStatus(res);

    if (status == PGRES_TUPLES_OK)
    {
        // 检查是否有查询结果
        if (PQntuples(res) > 0 && PQnfields(res) > 0)
        {
            // 检查字段是否为NULL
            if (!PQgetisnull(res, 0, 0))
            {
                const char *dirValue = PQgetvalue(res, 0, 0);
                size_t dirLen = strlen(dirValue);

                // 检查缓冲区大小是否足够
                if (dirLen + 1 <= bufferSize)
                { // +1 用于存储字符串结束符'\0'
                    strcpy(DataDirectory, dirValue);
                    result = 0; // 成功
                }
                else
                {
                    std::cerr << "错误：缓冲区太小，无法存储路径" << std::endl;
                    result = 5;
                }
            }
            else
            {
                std::cerr << "错误：查询结果为NULL" << std::endl;
            }
        }
        else
        {
            std::cerr << "错误：未找到data_directory配置" << std::endl;
        }
    }
    else
    {
        std::cerr << "查询执行失败：" << PQerrorMessage(conn) << std::endl;
    }

    // 释放结果集资源
    PQclear(res);
    return result;
}

int main(int argc, char **argv)
{
    const char *location;
    if (argc > 1)
        location = argv[1];
    else
        location = "threshold.yaml";
    std::string srouce;
    srouce = getsource(location);
    std::cout << srouce << std::endl;
    PGconn *conn = createPGConnection(srouce);
    if (!conn)
    {
        return EXIT_FAILURE;
    }
    YAML::Node config = YAML::LoadFile(location);
    std::string execute_sql;
    execute_sql = config["gaussdb"]["operate_type"]["terminate_sql"].as<std::string>();
    std::cout << execute_sql << std::endl;
    std::vector<std::string> pids;
    char dataDir[1024]; // 缓冲区，根据实际需求调整大小
    int ret = GetPGDataDirectory(conn, dataDir, sizeof(dataDir));
    if (ret == 0)
    {
        std::cout << "PostgreSQL数据目录路径:" << dataDir << std::endl;
    }
    else
    {
        std::cout << "获取数据目录失败，错误码：" << ret << std::endl;
    }
    double usage = -1.0;               // 初始化表示未获取
    check_disk_space(dataDir, &usage); // 传递指针
    // 使用获取到的使用率
    const size_t buffer_size = 1024;
    char cwd[buffer_size];
    get_current_working_directory(cwd, buffer_size);
    char path[1024]; // 缓冲区
    snprintf(path, sizeof(path), "%s/%s", cwd, "run.lock");
    if (usage >= 0)
    {
        if (usage >= config["gaussdb"]["disk_threshold"]["diskpercent"].as<std::int32_t>() && fileExists(path)==false)
        {
            printf("从main函数中获取的磁盘使用率: %.2f%%\n", usage);
            GetActivityPids(conn, execute_sql.c_str(), pids);
            for (size_t index = 0; index < pids.size(); ++index)
            {
                std::cout << index + 1 << ": " << pids[index] << std::endl;
                // 终止进程
                if (PGTerminateBackend(conn, pids[index]))
                {
                    std::cout << "成功终止进程 PID: " << pids[index] << std::endl;
                }
                else
                {
                    std::cout << "终止进程 PID: " << pids[index] << " 失败" << std::endl;
                }
            }
        ensureFileExists(path);
        }
    }
    else
    {
        printf("未能获取磁盘使用率\n");
    }

    PQfinish(conn);
    return 0;
}